package main

import (
	"bytes"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"sort"
	"strconv"
	"strings"
)

var httpHeader = http.Header{}

func main() {

	mapper := make(map[string]string)
	mapper["p"] = "port"
	mapper["h"] = "header"
	appArgs := createAppArgs(os.Args, mapper)

	httpHeaders := handleHeaderOption(appArgs)
	for key, val := range httpHeaders {
		httpHeader.Add(key, val)
	}

	var listen = ":8888"
	var filePath string

	portStr := appArgs.getOption("port")
	if portStr != "" {
		port, err := strconv.Atoi(portStr)
		if err == nil {
			listen = fmt.Sprintf("%s%d", ":", port)
		} else {
			listen = portStr
		}
	}

	if len(appArgs.args) != 0 {
		filePath = appArgs.args[0]
		// 查看是否是路径
		if !IsDir(filePath) {
			fmt.Printf("路径不存在或参数错误【%s】\n", filePath)
			return
		}
	}

	resData := appArgs.getOption("data")
	if resData != "" {
		http.HandleFunc("/data", func(writer http.ResponseWriter, request *http.Request) {
			handleResponseHeader(writer, httpHeader)
			fmt.Fprintf(writer, resData)
		})
	}

	fmt.Println("path: ", filePath)

	fmt.Println("监听：" + listen)
	if len(filePath) == 0 {
		http.HandleFunc("/forward", requestForward)
		http.HandleFunc("/", requestLog) // index 为向 url发送请求时，调用的函数
	} else {
		http.Handle("/", http.FileServer(http.Dir(filePath)))
	}

	// 检查是否启动udp
	udpPort := appArgs.getOption("udp")
	if udpPort != "" {
		go ListenUdp(udpPort)
	}

	http.ListenAndServe(listen, nil)

	println("退出")
}

// ListenUdp /**
func ListenUdp(portStr string) {
	fmt.Println("[UDP] 启动服务：", portStr)
	port, err := strconv.Atoi(portStr)
	if err != nil {
		log.Fatal("[UDP] 端口错误，", err)
	}
	udpConn, err := net.ListenUDP("udp", &net.UDPAddr{
		IP:   net.IPv4(0, 0, 0, 0),
		Port: port,
	})

	if err != nil {
		log.Fatal("[UDP] 启动错误", err)
		return
	}

	// 循环读取消息
	for {
		var data [1024]byte
		n, addr, err := udpConn.ReadFromUDP(data[:])
		if err != nil {
			log.Fatalf("[UDP] 读取错误【%s】：%s", addr, err)
			break
		}
		func() {
			// 返回数据
			fmt.Printf("[UDP] 收到【%s】数据：%v", addr, string(data[:n]))
			str := RandStr(10)
			fmt.Printf("[UDP] 发送消息: 【%s】\n\n", str)
			str = "[udp] " + str + "\n"
			_, err := udpConn.WriteToUDP([]byte(str), addr)
			if err != nil {
				log.Fatal("[UDP] 发送消息错误:", err)
			}
		}()
	}

}

func RandStr(n int) string {
	result := make([]byte, n/2)
	rand.Read(result)
	return hex.EncodeToString(result)
}

func requestLog(res http.ResponseWriter, req *http.Request) {
	if req.RequestURI == "/favicon.ico" {
		return
	}

	handleResponseHeader(res, httpHeader)

	var bt bytes.Buffer

	bt.WriteString("===============请求头================\n")
	bt = ws(bt, req.Method, " ", req.RequestURI, " ", req.Proto, "\n")

	// 计算key长度
	var length = 0
	var keys []string
	for k, _ := range req.Header {
		keys = append(keys, k)
		l := len(k)
		if l > length {
			length = l
		}
	}
	length += 5
	sort.Strings(keys)

	bt = ws(bt, strFmt("host:", length), req.Host, "\n")
	header := req.Header
	for _, k := range keys {
		bt = ws(bt, strFmt(k, length), ts(header[k]), "\n")
	}

	bt.WriteString("===============QueryString================\n")
	bt.WriteString(req.URL.RawQuery)
	bt.WriteString("\n")

	bt.WriteString("===============请求体================\n")
	body, _ := ioutil.ReadAll(req.Body)
	bt.WriteString(string(body))

	bt.WriteString("\n")

	str := bt.String()
	fmt.Println(str)
	fmt.Fprintf(res, str)
}

func requestForward(res http.ResponseWriter, req *http.Request) {
	fmt.Println("请求转发：")
	queryParam := req.URL.Query()
	urlParam := queryParam.Get("url")
	if urlParam == "" {
		res.WriteHeader(500)
		fmt.Fprintln(res, "请设置转发的url参数")
		return
	}

	// 发送请求
	fres, err := http.Get(urlParam)
	if err != nil {
		res.WriteHeader(500)
		fmt.Fprintln(res, "请求失败：", err.Error())
		return
	}
	defer fres.Body.Close()
	body, err := ioutil.ReadAll(fres.Body)
	if err != nil {
		res.WriteHeader(500)
		fmt.Fprintln(res, "请求错误：", err.Error())
		return
	}

	bodyStr := string(body)

	if fres.StatusCode >= 500 {
		res.WriteHeader(500)
		fmt.Fprintln(res, "请求错误：", fres.StatusCode, "BODY: ", bodyStr)
		return
	}

	fmt.Fprintln(res, bodyStr)
	fmt.Print(bodyStr)
}

func handleHeaderOption(args *AppArgs) map[string]string {
	m := make(map[string]string)

	arr := args.Options["header"]
	if arr == nil {
		return m
	}

	for _, item := range arr {
		keyVal := strings.SplitN(item, ":", 2)
		m[keyVal[0]] = keyVal[1]
	}

	return m
}

/*
*
处理响应头
*/
func handleResponseHeader(res http.ResponseWriter, header http.Header) {
	if header == nil {
		return
	}
	currHeader := res.Header()
	for key := range header {
		val := header.Get(key)
		currHeader.Add(key, val)
	}
	setHeader(currHeader, "Access-Control-Allow-Origin", "*")
	setHeader(currHeader, "Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE")
	setHeader(currHeader, "Access-Control-Allow-Headers", "Content-Type, Authorization")
}

func setHeader(header http.Header, key string, val string) {
	// 处理跨域
	h1 := header.Get(key)
	if h1 == "" {
		header.Add(key, val)
	}
}

// IsDir 判断目录是否存在
func IsDir(fileAddr string) bool {
	s, err := os.Stat(fileAddr)
	if err != nil {
		log.Println(err)
		return false
	}
	return s.IsDir()
}

func ts(arr []string) string {
	if arr == nil || len(arr) == 0 {
		return ""
	}
	if len(arr) == 1 {
		return arr[0]
	}
	s, _ := json.Marshal(arr)
	return string(s)
}

func ws(bt bytes.Buffer, str ...string) bytes.Buffer {
	for i := 0; i < len(str); i++ {
		bt.WriteString(str[i])
	}
	return bt
}

func strFmt(str string, length int) string {
	return fmt.Sprintf("%s%"+strconv.Itoa(length-len(str))+"s", str, "")
}
